import numpy as np
import json
import triton_python_backend_utils as pb_utils


class TritonPythonModel:

    def initialize(self, args):
        self.logger = pb_utils.Logger

        # parse model configs
        self.model_config = json.loads(args['model_config'])

        # parse parameters
        self.parameters = self.model_config['parameters']
        self.model_name = self.model_config['name']

        self.net_h   = float(self.parameters['net_h']['string_value'])
        self.net_w   = float(self.parameters['net_w']['string_value'])
        self.logger.log_info('net_h: {}'.format(self.net_h))
        self.logger.log_info('net_w: {}'.format(self.net_w))

        self.logger.log_info('{} started'.format(self.model_name))


    def execute(self, requests):
        print('running {}'.format(self.model_name))
        # print(type(requests)) #List
        responses = []
        for request in requests:
            # print(type(request))
            org_point_coords = pb_utils.get_input_tensor_by_name(request, 'org_point_coords').as_numpy()
            org_img_hw = pb_utils.get_input_tensor_by_name(request, 'org_img_hw').as_numpy()
            batch = org_img_hw.shape[0]
            # print(org_point_coords.shape)
            # print(type(org_point_coords))
            # print(type(org_img_hw))

            point_coords = []
            for i in range(batch):
                org_point_coord = org_point_coords[i, ...]
                org_point_coord[..., 0] = org_point_coord[..., 0] / org_img_hw[i, 1] * self.net_w
                org_point_coord[..., 1] = org_point_coord[..., 1] / org_img_hw[i, 0] * self.net_h
                org_point_coord = np.expand_dims(org_point_coord.astype(np.int32).astype(np.float32), axis=0)
                point_coords.append(org_point_coord)

            point_coords = np.vstack(point_coords)
            # print(point_coords.shape)

            inference_response = pb_utils.InferenceResponse(
                output_tensors=[
                    pb_utils.Tensor(
                        "point_coords",
                        point_coords,
                    ),
                ]
            )
            responses.append(inference_response)

        return responses


    def finalize(self):
        self.logger.log_info('Finalizing {}'.format(self.model_name))
        